﻿#include "qdiskmonitor.h"
#include <QDebug>

#ifdef Q_OS_WIN
#include <windows.h>
#else
#include <sys/statvfs.h>
#include <QFile>
#include <QTextStream>
#endif

QDiskMonitor::QDiskMonitor() {}
QDiskMonitor::~QDiskMonitor() {}

QList<DiskUseInfo> QDiskMonitor::getAllDiskUseInfo() {
#ifdef Q_OS_WIN
    return getDiskUseInfoWindows();
#else
    return getDiskUseInfoLinux();
#endif
}

#ifdef Q_OS_WIN
QList<DiskUseInfo> QDiskMonitor::getDiskUseInfoWindows() {
    QList<DiskUseInfo> diskInfos;

    DWORD drives = GetLogicalDrives();
    char drivePath[4] = "A:\\";

    for (char i = 'A'; i <= 'Z'; ++i) {
        if (drives & (1 << (i - 'A'))) { // Check if the drive is available
            drivePath[0] = i;
            ULARGE_INTEGER freeBytesAvailable, totalNumberOfBytes, totalNumberOfFreeBytes;
            if (GetDiskFreeSpaceExA(drivePath,
                                    &freeBytesAvailable,
                                    &totalNumberOfBytes,
                                    &totalNumberOfFreeBytes)) {
                DiskUseInfo info;
                info.driveLetter = QString::fromLatin1(drivePath);
                info.totalSize = totalNumberOfBytes.QuadPart;
                info.usageSize = (totalNumberOfBytes.QuadPart - totalNumberOfFreeBytes.QuadPart);

                // Get drive type
                UINT driveType = GetDriveTypeA(drivePath);
                switch (driveType) {
                    case DRIVE_UNKNOWN:      info.driveType = "Unknown"; break;
                    case DRIVE_NO_ROOT_DIR:  info.driveType = "No Root Directory"; break;
                    case DRIVE_REMOVABLE:    info.driveType = "Removable Disk"; break;
                    case DRIVE_FIXED:        info.driveType = "Local Disk"; break;
                    case DRIVE_REMOTE:       info.driveType = "Network Drive"; break;
                    case DRIVE_CDROM:        info.driveType = "CD-ROM"; break;
                    case DRIVE_RAMDISK:      info.driveType = "RAM Disk"; break;
                    default:                 info.driveType = "Unknown"; break;
                }

                // Get file system type
                char fileSystemName[MAX_PATH];
                if (GetVolumeInformationA(drivePath, nullptr, 0, nullptr, nullptr, nullptr, fileSystemName, MAX_PATH)) {
                    info.fileSystem = QString::fromLatin1(fileSystemName);
                } else {
                    info.fileSystem = "Unknown";
                }

                diskInfos.append(info);
            }
        }
    }

    return diskInfos;
}
#else
QList<DiskUseInfo> QDiskMonitor::getDiskUseInfoLinux() {
    QList<DiskUseInfo> diskInfos;

    QFile mountsFile("/proc/mounts");
    if (!mountsFile.open(QIODevice::ReadOnly | QIODevice::Text)) {
        qDebug() << "Failed to open /proc/mounts";
        return diskInfos;
    }

    QTextStream in(&mountsFile);
    QString line;
    while (!in.atEnd()) {
        line = in.readLine();
        QStringList parts = line.split(' ');
        if (parts.size() < 3) {
            continue; // Ignore incomplete lines
        }

        QString device = parts[0]; // First column is the device name
        QString mountPoint = parts[1]; // Second column is the mount point
        QString fileSystem = parts[2]; // Third column is the file system type

        struct statvfs vfs;
        if (statvfs(mountPoint.toUtf8().constData(), &vfs) == 0) {
            uint64_t totalSize = (static_cast<uint64_t>(vfs.f_blocks) * vfs.f_frsize);
            uint64_t freeSize = (static_cast<uint64_t>(vfs.f_bfree) * vfs.f_frsize);
            uint64_t usageSize = totalSize - freeSize;

            DiskUseInfo info;
            info.driveLetter = mountPoint;
            info.fileSystem = fileSystem;
            info.driveType = "Local Disk"; // On Linux, all are treated as local disks
            info.totalSize = totalSize;
            info.usageSize = usageSize;

            diskInfos.append(info);
        }
    }

    mountsFile.close();
    return diskInfos;
}
#endif
